# 画每台机器的异常分数及bestf1score确定的阈值
from matplotlib import colors, pyplot as plt
import pickle
import numpy as np
from sklearn.metrics import precision_recall_curve
from typing import Sequence, Tuple, Dict, Optional
import pathlib
import os
from global_data_2 import *

index_num = index_num
root_path = pathlib.Path(os.path.join(project_path, f"ctf_code/{exp_key}"))
# 存储各个簇选出的best f1 score对应的阈值
threshold_path = pathlib.Path(os.path.join(project_path, f"ctf_code/{exp_key}/label_data/threshold.npy"))

def _best_f1score(labels: np.ndarray, scores: np.ndarray) -> Tuple[float, float, float, float]:
    try:
        precision, recall, thresholds = precision_recall_curve(y_true=labels, probas_pred=scores)
    except:
        print(set(labels.tolist()))
        print(scores)
        raise Exception
    f1score = 2 * precision * recall / np.clip(precision + recall, a_min=-1000, a_max=None)

    best_threshold = thresholds[np.argmax(f1score)]
    best_precision = precision[np.argmax(f1score)]
    best_recall = recall[np.argmax(f1score)]

    return best_threshold, best_precision, best_recall, np.max(f1score)

deviation_num = 1

def plot_score(label_list, value_list, score_list, score_all, reconstruct_data, reconstruct_mean, reconstruct_std, threshold, save_path):
    # 前19个子图画value，最后一个图画score和阈值
    fig = plt.figure(figsize=(20, 40))

    ax_score = fig.add_subplot(index_num + 1, 1, 1)
    ax_score.plot(range(len(label_list)), score_list, color='green')
    # ax_score.set_xticks([])
    # ax_score.set_yticks([])
    if np.sum(label_list) > 0:
        anomaly_list = np.copy(score_list)
        anomaly_list[np.where(label_list==0)[0]] = np.inf
        ax_score.plot(range(label_list.shape[0]), anomaly_list, color='red', linewidth=3)
    ax_score.hlines(threshold, 0, label_list.shape[0] - 1, colors='black', linewidth=3)

    for row in range(index_num):
        ax = fig.add_subplot(index_num + 1, 1, index_num + 1 - row)
        ax.set_ylabel(f"{row}")
        ax_score_a = ax.twinx()
        ax.plot(range(label_list.shape[0]), value_list[omni_windows_size - 1:, row], color='grey')
        # ax.plot(range(label_list.shape[0]), reconstruct_data[:, row], color='purple')
        ax.plot(range(label_list.shape[0]), reconstruct_mean[:, row], color='blue')
        ax.plot(range(label_list.shape[0]), reconstruct_mean[:, row] - deviation_num * reconstruct_std[:, row], color='yellow')
        ax.plot(range(label_list.shape[0]), reconstruct_mean[:, row] + deviation_num * reconstruct_std[:, row], color='orange')
        # 每一个维度单独的异常分数
        ax_score_a.plot(range(label_list.shape[0]), score_all[:, row], color='green')
        if np.sum(label_list) > 0:
            anomaly_value_list = np.copy(value_list[omni_windows_size - 1:, row])
            anomaly_value_list[np.where(label_list==0)[0]] = np.inf
            ax.plot(range(label_list.shape[0]), anomaly_value_list, color='red', linewidth=3)
            
        # if np.sum(label_list) > 0:
        #     ax.axvspan(np.where(label_list == 1)[0][0], np.where(label_list == 1)[0][-1], color="red", alpha=0.3,
        #             label='Anomaly')
        ax.set_xticks([])
        # ax.set_yticks([])
    plt.subplots_adjust(wspace=0, hspace=0.2)
    plt.savefig(save_path, bbox_inches="tight")
    plt.close(fig)



if __name__ == '__main__':
    # 获取每个族的所有机器的检测结果，按照一定顺序拼接起来
    with open(root_path / 'z_results/z_cluster.pkl', mode='rb') as f:
        cluster_list = pickle.load(f)
    threshold_list = np.load(threshold_path)
    # 读取历史数据的均值和方差，对重构数据进行逆标准化
    with open(os.path.join(project_path, f'result/{exp_key}/historical_data/data_info.pkl'), 'rb') as f:
            [mean_matrix, std_matrix] = pickle.load(f)
            
    # 获取每个簇的机器列表
    for cluster_index, machine_list in enumerate(cluster_list):
        # 读取异常分数
        threshold = threshold_list[cluster_index]
        for machine_index, machine_id in enumerate(machine_list):
            mean = mean_matrix[int(machine_id), :]
            std = std_matrix[int(machine_id), :]
            # machine_index用于寻找对应机器的异常标注数据
            with open(root_path / f'label_data/{cluster_index + 1}_score.pkl', mode='rb') as f:
                score_list = pickle.load(f)[0][:, machine_index]
            score_all = np.load(root_path / f'label_data/{cluster_index + 1}_score_all.npy')[:,machine_index, :]
            with open(os.path.join(project_path, f'label_result/{int(machine_id)}.pkl'), mode='rb') as l:
                label_list = pickle.load(l)[omni_windows_size - 1:]
            value_list = np.genfromtxt(os.path.join(project_path, f'CTF_data/{int(machine_id)}_{6}.txt'), delimiter=',')
            for day in range(7, 8):
                value_list = np.vstack(
                    [value_list, np.genfromtxt(os.path.join(project_path, f'CTF_data/{int(machine_id)}_{day}.txt'), delimiter=',')])
            # 读取重构数据列表, 逆标准化
            reconstruct_data_cluster = np.load(os.path.join(project_path, f"result/{exp_key}/result_for_period2_{cluster_index+1}/reconstruct_data.npy"))
            reconstruct_data = reconstruct_data_cluster[machine_index::len(machine_list), :]
            reconstruct_data = reconstruct_data * (std + 1e-9) + mean
            reconstruct_mean_cluster = np.load(os.path.join(project_path, f"result/{exp_key}/result_for_period2_{cluster_index+1}/reconstruct_mean.npy"))
            reconstruct_mean = reconstruct_mean_cluster[machine_index::len(machine_list), :]
            reconstruct_mean = reconstruct_mean * (std + 1e-9) + mean
            reconstruct_std_cluster = np.load(os.path.join(project_path, f"result/{exp_key}/result_for_period2_{cluster_index+1}/reconstruct_std.npy"))
            reconstruct_std = reconstruct_std_cluster[machine_index::len(machine_list), :]
            reconstruct_std = reconstruct_std * (std + 1e-9) + mean

            if len(score_list) != len(label_list):
                print(f"len(score_list):{len(score_list)}, len(label_list):{len(label_list)}")
                raise Exception
            save_path = pathlib.Path(os.path.join(project_path, f'pic/{exp_key}')) / str(cluster_index) / f"{int(machine_id)}.png"
            save_path.parent.mkdir(parents=True, exist_ok=True)
            plot_score(np.array(label_list), value_list, -np.array(score_list), -score_all, reconstruct_data, reconstruct_mean, reconstruct_std, threshold, save_path)
            print(f"{cluster_index}->{machine_id}")
            # raise Exception
